--------------------------------------------
--  Phaedon Sinis (sinis@stanford.edu)
--  CS240h, Lab 1
--  Fall 2011
--------------------------------------------


import Data.Char
import Data.List (sortBy)
import qualified Data.Map as Map
import System.Environment


----------------------
--  String operations
----------------------

-- Convert a string to lower case
toLowerWord :: String -> String
toLowerWord = map toLower


-- catches possessives ("David's") and contractions ("won't")
isPossessive :: String -> Bool
isPossessive (a:(b:(c:xs))) = isAlpha a
                      && b == '\''
                      && (c == 's') || (c == 't')
                      && isSeparator (head xs)
isPossessive _ = False


-- catches Irish names like "O'Sullivan"
isIrish :: String -> Bool
isIrish (a:(b:xs)) = a == 'O'
                 && b == '\''
                 && isAlpha (head xs)
isIrish _ = False
              

-- We want to keep apostrophes, but not single-quotes. Get 'em out!!
removeSingleQuotes :: String -> String
removeSingleQuotes (x:xs) 
  | isPossessive (x:xs) 
    || isIrish (x:xs) = x : head xs : removeSingleQuotes (tail xs)
  | x == '\'' = removeSingleQuotes xs                                          
  | otherwise = x : removeSingleQuotes xs
removeSingleQuotes [] = ""


-- Filter punctuation out of a string
noPunctuation :: String -> String
noPunctuation str = filter isNotPunctuation (removeSingleQuotes str)
  where isNotPunctuation s = isAlphaNum s || isSpace s || s == '\''
 

-- Split a string into all lower-case words, excluding punctuation
splitWords :: String -> [String]
splitWords str = words (toLowerWord (noPunctuation str))


----------------------------
--  Data structure manip
----------------------------

-- Takes a list of words and creates a map that tallies up the word counts
wordCountMap :: [String] -> Map.Map String Int
wordCountMap [] = Map.empty
wordCountMap (x:xs) = Map.insertWith (+) x 1  $ wordCountMap xs


-- Takes a list of pairs and flips 'em around!
swapKeyValPairs :: [(a, b)] -> [(b, a)]
swapKeyValPairs kvList = [(v, k) | (k, v) <- kvList]


-- Generates a single line for our final output
formatLine ::  Int -> (Int, String) -> String
formatLine len (i, s) 
  | i == 0 = "" -- ignore if no bars
  | otherwise = s ++ spaces ++ hashes ++ "\n"
    where
      hashes = replicate i '#' 
      spaces = replicate (len - length s) ' '


-- Generates our final output!
formatOutput :: [(Int, String)] -> Int -> String
formatOutput (x:xs) len 
  = formatLine len x ++ formatOutput xs len
formatOutput [] _ = ""


-- Scales the word counts by a constant factor
scalePairs :: [(Int, String)] -> Float -> [(Int, String)]
scalePairs pairList scaleFactor 
  | scaleFactor >= 0.999 = pairList -- essentially, 1 (don't scale)
  | otherwise = map scaleInt pairList
    where
      scaleInt (i, s) = (floor (fromIntegral i * scaleFactor), s)


-- Comparator for sorting (Int, String) pairs:
--  if freqs are different, sort in descending order (numerical)
--   if they are the same, sort ascending (alphabetical)
myCmp :: (Ord a, Ord b) => (a, b) -> (a, b) -> Ordering
myCmp a b | fst a < fst b                   = GT
          | fst a == fst b && snd a > snd b = GT
          | otherwise                       = LT


-- Reverse keys and values in a map, and return a 
-- sorted list of the (value, key) pairs
sortedRevMapList :: Map.Map String Int 
                    -> ( (Int, String) -> (Int, String) -> Ordering )
                    ->  [(Int, String)]
sortedRevMapList wordMap cmp = vkSorted
  where
    -- change to a list:
    ascKeyList = Map.toAscList wordMap
    
    -- swap keys & values
    vkList = swapKeyValPairs ascKeyList
    
    -- and sort it
    vkSorted = sortBy cmp vkList


-- compute the length of the longest word 
-- (so that we know the column num for the '#' symbols)
maxWordLen :: Map.Map String Int -> Int
maxWordLen wordMap = w + 1
  where
    lenMap = Map.mapWithKey (\k _ -> length k) wordMap
    vkSorted = sortedRevMapList lenMap myCmp
    w = fst (head vkSorted)


-- Computes a scaling factor in [0, 1] based on the longest line.
-- This factor will be used to linearly scale all of the bars, 
-- so that no output column is longer than 80 chars.
histoScaleFactor :: Int -> Int -> Float
histoScaleFactor maxFreq biggestWord | rawLineLen <= 80 = 1.0
                                     | otherwise = n / d
  where 
    rawLineLen = maxFreq + biggestWord + 1
    n = 80.0 - 1.0 - fromIntegral biggestWord
    d = fromIntegral maxFreq


-- Converts a word map into a formatted text histogram (final output!!) 
histo :: Map.Map String Int -> String
histo wordMap = formatOutput (sortBy myCmp vkScaled) biggestWord
  where 
    -- Generate a sorted list of pairs from the word map
    vkPreSorted = sortedRevMapList wordMap myCmp

    biggestWord = maxWordLen wordMap
      
    -- in order to (possibly) scale the hash marks, 
    -- we need the longest freq length:
    maxFreq = fst (head vkPreSorted)

    histScale = histoScaleFactor maxFreq biggestWord
    
    -- Why do we sort again? Because in scaling the frequencies,
    -- many frequencies may round to the same scaled frequency,
    -- thereby ruining our nice sorting job from before
    vkScaled = scalePairs vkPreSorted histScale

      
----------------------------
--  The program!
----------------------------
      
main :: IO ()
main = do     
  args <- getArgs
  
  -- no args? retrieve from stdin:
  if null args then 
    do putStr "Please enter anything you like. Press CTRL-D to finish.\n"
       text <- getContents
       putStr $ histo (wordCountMap (splitWords text))
       
  -- args? then treat the first one as a filename and retrieve the contents:
    else
    do text <- readFile (head args)
       putStr $ histo (wordCountMap (splitWords text))

